/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package no.ffi.tpm.aik;

import iaik.tc.apps.jtt.aik.Client;
import iaik.tc.apps.jtt.aik.Constants;
import iaik.tc.apps.jtt.common.CommonSettings;
import iaik.tc.jsr321.tpm.TPMContextImpl;
import iaik.tc.jsr321.tpm.keys.TPMKeyImpl;
import iaik.tc.tss.api.constants.tpm.TcTpmConstants;
import iaik.tc.tss.api.constants.tsp.TcTssConstants;
import iaik.tc.tss.api.structs.common.TcBlobData;
import iaik.tc.tss.api.structs.tpm.TcTpmKeyParms;
import iaik.tc.tss.api.structs.tpm.TcTpmPubkey;
import iaik.tc.tss.api.structs.tpm.TcTpmQuoteInfo2;
import iaik.tc.tss.api.structs.tpm.TcTpmSymCaAttestation;
import iaik.tc.tss.api.structs.tpm.TcTpmSymmetricKeyParms;
import iaik.tc.tss.api.structs.tsp.TcTssUuid;
import iaik.tc.tss.api.structs.tsp.TcUuidFactory;
import iaik.tc.tss.api.tspi.TcIContext;
import iaik.tc.tss.api.tspi.TcIPolicy;
import iaik.tc.tss.api.tspi.TcIRsaKey;
import iaik.tc.tss.api.tspi.TcITpm;
import iaik.tc.tss.impl.csp.TcCrypto;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.security.PublicKey;
import java.security.cert.X509Certificate;
import java.security.interfaces.RSAPublicKey;
import java.util.Arrays;
import java.util.UUID;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.trustedcomputing.tpm.TPM;
import javax.trustedcomputing.tpm.TPMContext;
import javax.trustedcomputing.tpm.keys.IdentityKey;
import javax.trustedcomputing.tpm.keys.KeyManager;
import javax.trustedcomputing.tpm.keys.StorageRootKey;
import javax.trustedcomputing.tpm.structures.PCRInfo;
import javax.trustedcomputing.tpm.structures.Secret;
import javax.trustedcomputing.tpm.structures.ValidationData;
import javax.trustedcomputing.tpm.tools.Attestor;
import no.ffi.gismoidm.certstore.CertificateFileOperations;
import no.ffi.tpm.utils.CaServiceTest;
import no.ffi.tpm.crypto.*;

import no.ffi.tpm.TPMUtils;


/**
 *
 * @author Federico Mancinie <fma at ffi.no>
 */
public class AIKUtils {
    public static String uuid;
    public final static String REQUEST="Req";
    public final static String RESPONSE="Resp";
    public static ValidationData val;
     /**
     * This method takes the output of a collateIdentityRequest before is encrypted 
     * in the TSS and transform it in a AIKRequest object before to encrypt it with
     * the PCA public key, generating an EncryptedObject.
     * @param PCAKey
     * @return 
     */    
    public static EncryptedObject createAIKRequest(String CACertFile, String AIKPass, String TPMPass) {
       
        EncryptedObject aikReqByte = null;
       
        try { 
            //Get the PCA public key from its certificate
            System.out.println(System.getProperty("user.dir"));
            X509Certificate cert = CertificateFileOperations.loadPemCert(CACertFile);
            RSAPublicKey PCAKey = (RSAPublicKey) cert.getPublicKey();
            //Exploit java_tpm_tools to create a collateIdentityRequest
            Client client = new Client(CommonSettings.getTssFactory());
            TcBlobData pcaPubKey=TcCrypto.pubJavaToTpmKey(PCAKey).getEncoded();
            TcBlobData pcaLabel=TcBlobData.newString("PCALabel", false);;
            TcBlobData srkSecr=TcBlobData.newByteArray(TcTssConstants.TSS_WELL_KNOWN_SECRET);
            long srkMode=TcTssConstants.TSS_SECRET_MODE_SHA1;        
            TcBlobData aiksecret= TcBlobData.newString(AIKPass, false);
            TcBlobData ownerSecret=TcBlobData.newString(TPMPass, false);
            //Get the Endorsment certificate used to build the request
            TcBlobData ek=TPMUtils.getEK(ownerSecret);
            client.overrideEkCertificate(ek);
            System.out.println(new String(pcaPubKey.toHexString()) + "SRK="+srkSecr.toHexString()+" ownerSecret= "+ownerSecret.toString()+ " aikpass = "+aiksecret.toString());
            TcBlobData TPMRequest=client.collateIdentityReq2(pcaPubKey, srkSecr, srkMode, ownerSecret, aiksecret, pcaLabel);
            
            //Get AIK and store it
            TcIRsaKey aik = client.getAikKey();
            TPMContext cont = TPMContext.getInstance(); 
            cont.connect(null);
            TcIContext context = ((TPMContextImpl) cont).getTcIContext();//CommonSettings.getTssFactory().newContextObject();
            //context.connect(CommonSettings.getHostname());
            
            //context.getTpmObject().quote(aik, null, null);
            TcTssUuid srkUuid=TcUuidFactory.getInstance().getUuidSRK();
            TcTssUuid aikUuid=TcUuidFactory.getInstance().generateRandomUuid();
            long storageType=TcTssConstants.TSS_PS_TYPE_SYSTEM;
            context.registerKey(aik,storageType , aikUuid, storageType, srkUuid);
            uuid=aikUuid.toStringNoPrefix();
            System.out.println("AIK register with uuid "+uuid);
            //Do a quote with the newly created AIK
            KeyManager manager = cont.getKeyManager();
            StorageRootKey srk = manager.loadStorageRootKey(Secret.WELL_KNOWN_SECRET);
            IdentityKey aik2 =(IdentityKey) manager.loadTPMSystemKey(srk, UUID.fromString(uuid), cont.getSecret(AIKPass.toCharArray()));
            Attestor att=cont.getAttestor();
            int[] pcrIdx={0,1,2,3,4,5,6,7,8,9,10,11,12};
            //System.out.println(Arrays.toString(pcrIdx));
            TPM tpm = cont.getTPMInstance();
            PCRInfo pcrs = tpm.readPCR(pcrIdx);
            
            val=att.quote(pcrIdx, aik2, cont.getDigest(new byte[20]));
            System.out.println(val.toString());
            /*//Check whether AIK can be loaded
            TPMContext context2 = TPMContext.getInstance(); 
            context2.connect(null);
            KeyManager manager = context2.getKeyManager();
            StorageRootKey srk = manager.loadStorageRootKey(Secret.WELL_KNOWN_SECRET);
            IdentityKey aik2 =(IdentityKey) manager.loadTPMSystemKey(srk, UUID.fromString(aikUuid.toStringNoPrefix()), context2.getSecret("AIKSecret".toCharArray()));
            */
            //parse request with custom class
            //TPMReq req=new TPMReq(TPMRequest.asByteArray());//decode request
            //byte[] encrSymmKey=req.getAsymBlob();//The session key encrypted with PCA public key
            //byte[] encrReqWithIv=req.getSymBlobwithIV();
            //Decrypt the request, parse it and make an AIKRequest object out of it
            //Create an ecnrypted object so we can reuse existing code
            //EncryptedObject encObj=new EncryptedObject(CryptoOperations.asymmetricAlgPCA,8,encrReqWithIv,encrSymmKey,null);
            //byte[] clearTextTPMReq=(byte[]) crypto.decrypt(encObj);//The request generated by the TPM
            AIKRequest  aikReq=PCAReqParser.createAIKRequest(TPMRequest.asByteArray(),PCAKey);  
            aikReq.setPCR(val.getData(),val.getValidationData(),val.getNonce());
            PublicKey aikp = aikReq.getAIK();
            //System.out.println(aikReq.getEC().toString());
            //Encrypt the AIKRequest with the actual PCA public key as an EncryptedObject
            aikReqByte=CryptoOperations.encrypt(aikReq, PCAKey,CryptoOperations.asymmetricAlgPCA);           
        } catch (Exception ex) {
            Logger.getLogger(AIKUtils.class.getName()).log(Level.SEVERE, null, ex);
        }
         return aikReqByte;
    }
    
    /**
     * Sends the AIKRequest to the PCA/IdP and returns the response
     * @param PCAReq The EncryptedObject containing the AIKRequest object
     * @return The EncryptedObject containing the PCA response
     * @throws Exception 
     */
    public static EncryptedObject certifyAIK(EncryptedObject PCAReq) throws Exception{
        Object responseObj = CaServiceTest.toPCA(PCAReq,CaServiceTest.serviceURLAIK);
        EncryptedObject resp=(EncryptedObject) responseObj;
        return resp;
    }
    
   
    /**
     * Activate an AIK using the response from the PCA and returns the secret 
     * generated by the PCA and encrypted with the session key in turn encrypted
     * with the EK
     * @param resp The EncryptedObject containing the PCA response
     * @param TPMPass The TPM password
     * @return The secret generated by the PCA that needs to be used in conjunction with 
     * the AIK. Normally this would be an AIK certificate
     */
    public static String activateAIK(EncryptedObject resp,String TPMPass,String AIKpass) {
        String secret=null;
        try {
            Object[] PCAResp=parseResponse(resp);
            //Use JSR321 to load the AIK
            //Connect to the TPM and create a context
            TPMContext context = TPMContext.getInstance();
            context.connect(null);
            //Load the AIK stored in the system storage by UUID
            KeyManager manager = context.getKeyManager();
            StorageRootKey srk = manager.loadStorageRootKey(Secret.WELL_KNOWN_SECRET);
            IdentityKey aik2 = (IdentityKey) manager.loadTPMSystemKey(srk, UUID.fromString(uuid), context.getSecret(AIKpass.toCharArray()));
            System.out.println(" aik modulus="+ aik2.getPublicKey().getModulus());
            //Retrieve the actual AIK key stored in a tpm compatible format
            //This requires using the underlying reference implementation of the JSR321
            TcIRsaKey aikKey = ((TPMKeyImpl) aik2).getTciRsaKey();
            //Retrieve also the actual tpm context from the reference implementation
            //so to have access to the internal methods
            TcIContext context2 = ((TPMContextImpl) context).getTcIContext();
            //Retrieve the TPM object
            TcITpm tpm = context2.getTpmObject();
            //Set the TPM owner password to not get a pop-up window asking for it
            TcIPolicy tpmPolicy = tpm.getPolicyObject(TcTssConstants.TSS_POLICY_USAGE);
            tpmPolicy.setSecret(TcTssConstants.TSS_SECRET_MODE_PLAIN, TcBlobData.newString(TPMPass));
            //Call the activateIdentity method that decrypts the session key inside the
            //TPM and then use the session key to decrypt the AIK certificate
            TcBlobData aikCert = tpm.activateIdentity(aikKey, (TcBlobData) PCAResp[1], (TcBlobData) PCAResp[0]);
            //Print the AIK certificate
            secret=new String(aikCert.asByteArray());
            
        } catch (Exception ex) {
            Logger.getLogger(AIKUtils.class.getName()).log(Level.SEVERE, null, ex);
        }
        return secret;
    }
    
    
        /**
     * This method creates the symmetricCABlob and asymmetricCABlob
     * as indicated in the TPM  specifications using the 
     * encryptedBuffer and encryptedKey objects in EncryptedObject containing 
     * the PCA response to an AIK certification request
     * @param eo The EncryptedObject containing the PCA response
     * @return the symmetricCABlob and asymmetricCABlob
     */
  static private Object[] parseResponse(EncryptedObject eo){
        byte[] asymBlob = eo.encryptedKey;
        byte[] encrBuffer = eo.encryptedBuffer;
        byte[] iv = Arrays.copyOf(encrBuffer, 16);
        byte[] symBlob = new byte[encrBuffer.length - 16];
        System.arraycopy(encrBuffer, 16, symBlob, 0, symBlob.length);
        // step 1: create the symmetric CA blob (TPM_SYM_CA_ATTESTATION) containing the
        // AIK credential

        TcTpmSymmetricKeyParms symParams = new TcTpmSymmetricKeyParms();
        symParams.setKeyLength(Constants.SYM_KEY_LEN);
        symParams.setBlockSize(Constants.SYM_BLOCK_SIZE);
        symParams.setIV(TcBlobData.newByteArray(iv));

        TcTpmKeyParms symCaAttestationKeyParms = new TcTpmKeyParms();
        symCaAttestationKeyParms.setAlgorithmID(Constants.SYM_ALGO_TPM);
        // TPM level structures are filled with TPM level constants
        symCaAttestationKeyParms.setEncScheme((int) TcTpmConstants.TPM_ES_SYM_CBC_PKCS5PAD);
        symCaAttestationKeyParms.setSigScheme((int) TcTpmConstants.TPM_SS_NONE);
        symCaAttestationKeyParms.setParms(symParams.getEncoded());

        TcTpmSymCaAttestation symCaAttestation = new TcTpmSymCaAttestation();
        symCaAttestation.setAlgorithm(symCaAttestationKeyParms);
        symCaAttestation.setCredential(TcBlobData.newByteArray(symBlob));
        TcBlobData symBlobEncrypted = symCaAttestation.getEncoded();


        // return the blob
        // symCaAttestation.credential holds the symmetrically encrypted AIK credential (the
        // rest of the structure is not encrypted).
        // asymCaContentsEncrypted holds symmetric session key (plus some key info) encrypted
        // with the public EK of the client.
        return new Object[]{symBlobEncrypted, TcBlobData.newByteArray(asymBlob)};
    }
    
  /**
     * This method can be used to store the PCA response to a file, so that 
     * the AIK can be activated at a later time
     * @param pcaReq 
     */
    public static void writeToFile(EncryptedObject eo,String type) throws Exception{
         if (!(type.equals(REQUEST)||type.equals(RESPONSE))) {
     Exception ex=new Exception("Type parameter must be of type AIKUtils.RESPONSE or AIKUtils.REQUEST");
                    throw  ex;    
 }
        byte[] pcaReqSerialized=CryptoOperations.serialize(eo);
            FileOutputStream fo;
            
        try {
            //Store the uuid of the AIK used in the request in the file name
            fo = new FileOutputStream("PCA"+type+"Serialized-"+uuid+".file");
            fo.write(pcaReqSerialized);
	    fo.close();
        } catch (IOException ex) {
            Logger.getLogger(AIKUtils.class.getName()).log(Level.SEVERE, null, ex);
        } 
	    
    }
    
    /**
     * Retrieves a stored AIKRequest object previously created for the AIK 
     * with the given UUID.
     * @param uuid
     * @return A serialized AIKRequest Object
     * @throws Exception 
     */
  public static Object getStoredAIKFiles(String uuid,String type) throws Exception{
 if (uuid == null || uuid.equals("")) {
     Exception ex=new Exception("uuid cannot be empty");
                    throw  ex;    
 }
 if (!(type.equals(REQUEST)||type.equals(RESPONSE))) {
     Exception ex=new Exception("Type parameter must be of type AIKUtils.RESPONSE or AIKUtils.REQUEST");
                    throw  ex;    
 }
    File folder = new File(".\\");
    File[] listOfFiles = folder.listFiles();
    File out=null;
    for (int i = 0; i < listOfFiles.length; i++){
        if (listOfFiles[i].isFile()){
                String fileName = listOfFiles[i].getName();
                System.out.println(fileName);
                if (fileName.startsWith("PCA"+type+"Serialized-")) {
                String fileUuid = fileName.substring(("PCA"+type+"Serialized-").length(), fileName.length() - ".file".length());
                if (uuid.equals(fileUuid)) {
                    out = listOfFiles[i];
                    break;
                }

            }
            }
        }  
    FileInputStream fi=new FileInputStream(out);
    byte[] req=new byte[fi.available()];
    fi.read(req);
    fi.close();
    
    return req;
  
  }
    


}

